#include <iostream>
#include <vector>
#include <ctime>
//#define VERBOSE

#include "augmentOpencv.h"
#include "KinectWrapper.h"

using namespace std;
using namespace cv;

const int WIDTH = 640;
const int HEIGHT = 480;

vector< vector<Point3f> > objPts;
vector< vector<Point2f> > imgCPts;
vector< vector<Point2f> > imgGPts;
int saveIdx = 0;

void saveCaliResult(string fn, vector< vector<Point3f> > objPts, Size size, vector< vector<Point2f> > imgPts, 
	vector<Mat> rvecs, vector<Mat> tvecs, Mat intriMat, Mat distortMat) 
{
	FileStorage fs(fn, FileStorage::WRITE);

	vector<float> flatVec;
	for (int i=0; i<objPts.size(); ++i) {
		for (int j=0; j<objPts[i].size(); ++j) {
			flatVec.push_back(objPts[i][j].x);
			flatVec.push_back(objPts[i][j].y);
			flatVec.push_back(objPts[i][j].z);
		}
		flatVec.push_back(-1);
		flatVec.push_back(-1);
		flatVec.push_back(-1);
	}
	Mat objPtsMat(flatVec.size()/3, 3, CV_32FC1, flatVec.data());
	fs << "objPtsMat" << objPtsMat;

	fs << "Size" << size;

	flatVec.clear();
	for (int i=0; i<imgPts.size(); ++i) {
		for (int j=0; j<imgPts[i].size(); ++j) {
			flatVec.push_back(imgPts[i][j].x);
			flatVec.push_back(imgPts[i][j].y);
			flatVec.push_back(1);
		}
		flatVec.push_back(-1);
		flatVec.push_back(-1);
		flatVec.push_back(-1);
	}
	Mat imgPtsMat(flatVec.size()/3, 3, CV_32FC1, flatVec.data());
	fs << "imgPtsMat" << imgPtsMat;

	fs << "rvecs" << "[";
	for (int i=0; i<rvecs.size(); ++i)
		fs << rvecs[i];
	fs << "]";

	fs << "tvecs" << "[";
	for (int i=0; i<tvecs.size(); ++i)
		fs << tvecs[i];
	fs << "]";

	fs << "intriMat" << intriMat;
	fs << "distortMat" << distortMat;
	fs.release();
}

void caliColor(void) {
	KinectDepthColorWrapper kdcWrapper;
	CvMat *depthRes = NULL, *colorRes = NULL, *mappedRes = NULL;
	Mat depth, color, mapped;

	kdcWrapper.init(WIDTH, HEIGHT);

	while (true) {
		kdcWrapper.retrieve(depthRes, colorRes, mappedRes);
		depth = Mat(depthRes);
		color = Mat(colorRes);
		mapped = Mat(mappedRes);

		imshow("depth", depth);
		imshow("color", color);
		imshow("mapped", mapped);
		char ch = cvWaitKey(50);
		if (ch == 27) break;
		if (ch == 13) {
			bool finish = false;

			for (int i=9; i>=5; --i) {
				if (finish) break;
				for (int j=i; j>=5; --j) {
					Size patternSize(i,j);
					vector<Point2f> corners;
					int res = findChessboardCorners(color, patternSize, corners, 1+8);
					cout << i << " " << j << " " << res << endl;
					if (res == 0) continue;

					Mat drawImg = color;
					drawChessboardCorners(drawImg, patternSize, corners, res);
					imshow("corner", drawImg);

					ch = cvWaitKey(0);

					if (ch == 13) {
						cout << "save it..." << endl;
						finish = true;
						char prefix[100];
						sprintf(prefix, "data%d", saveIdx++);
						string prefixStr(prefix);
						imwrite(prefixStr+"_c.jpg", color);
						FileStorage fs(prefixStr+"_depthValue.yml", FileStorage::WRITE);
						fs << "mapped" << mapped;
						fs << "depth" << depth;
						fs.release();
						imwrite(prefixStr+"_d.jpg", mapped);
						imwrite(prefixStr+"_depthOriginalImage.jpg", depth);

						vector<Point3f> objpts;

						for (int a=0; a<j; ++a) {
							for (int b=0; b<i; ++b) {
								objpts.push_back(Point3f(b,a,0));
							}
						}
						objPts.push_back(objpts);
						imgCPts.push_back(corners);

						break;
					}
				}
			}
		}
	}

	cout << "now calibration" << endl;
	Mat intriMat, distortMat;
	vector<Mat> rvecs, tvecs;
	calibrateCamera(objPts, imgCPts, Size(HEIGHT, WIDTH), intriMat, distortMat, rvecs, tvecs);
	ezPrintMat(intriMat);
	ezPrintMat(distortMat);

	saveCaliResult("calibrate_color_result.yml", objPts, Size(HEIGHT,WIDTH), imgCPts, rvecs, tvecs,
		intriMat, distortMat);
}

void caliGray(void) {
	VideoCapture capture(CV_CAP_DSHOW+1);
	capture.set(CV_CAP_PROP_FRAME_WIDTH, WIDTH);
	capture.set(CV_CAP_PROP_FRAME_HEIGHT, HEIGHT);

	while (true) {
		Mat gray;
		capture.read(gray);
		flip(gray, gray, 1);
		imshow("gray", gray);
		char ch = cvWaitKey(50);
		if (ch == 27) break;
		if (ch == 13) {
			bool finish = false;

			for (int i=9; i>=5; --i) {
				if (finish) break;
				for (int j=i; j>=5; --j) {
					Size patternSize(i,j);
					vector<Point2f> corners;
					int res = findChessboardCorners(gray, patternSize, corners, 1+8);
					cout << i << " " << j << " " << res << endl;
					if (res == 0) continue;

					Mat drawImg = gray;
					drawChessboardCorners(drawImg, patternSize, corners, res);
					imshow("corner", drawImg);

					ch = cvWaitKey(0);

					if (ch == 13) {
						cout << "save it..." << endl;
						finish = true;
						char prefix[100];
						sprintf(prefix, "data%d", saveIdx++);
						string prefixStr(prefix);
						imwrite(prefixStr+"_g.jpg", gray);

						vector<Point3f> objpts;

						for (int a=0; a<j; ++a) {
							for (int b=0; b<i; ++b) {
								objpts.push_back(Point3f(b,a,0));
							}
						}
						objPts.push_back(objpts);
						imgGPts.push_back(corners);

						break;
					}
				}
			}
		}
	}

	cout << "now calibration" << endl;
	Mat intriMat, distortMat;
	vector<Mat> rvecs, tvecs;
	calibrateCamera(objPts, imgGPts, Size(HEIGHT, WIDTH), intriMat, distortMat, rvecs, tvecs);
	ezPrintMat(intriMat);
	ezPrintMat(distortMat);

	saveCaliResult("calibrate_gray_result.yml", objPts, Size(HEIGHT,WIDTH), imgGPts, rvecs, tvecs,
		intriMat, distortMat);
}


void caliBoth(void) {
	KinectDepthColorWrapper kdcWrapper;
	CvMat *depthRes = NULL, *colorRes = NULL, *mappedRes = NULL;
	Mat depth, color, mapped, gray;

	kdcWrapper.init(WIDTH, HEIGHT);

	VideoCapture capture(CV_CAP_DSHOW+1);
	capture.set(CV_CAP_PROP_FRAME_WIDTH, WIDTH);
	capture.set(CV_CAP_PROP_FRAME_HEIGHT, HEIGHT);



	while (true) {
		kdcWrapper.retrieve(depthRes, colorRes, mappedRes);
		capture.read(gray);
		flip(gray, gray, 1);

		depth = Mat(depthRes);
		color = Mat(colorRes);
		mapped = Mat(mappedRes);

		imshow("depth", depth);
		imshow("color", color);
		imshow("mapped", mapped);
		imshow("gray", gray);

		char ch = cvWaitKey(50);
		if (ch == 27) break;

		if (ch == 13) {
			bool finish = false;

			for (int i=9; i>=5; --i) {
				if (finish) break;
				for (int j=i; j>=5; --j) {
					Size patternSize(i,j);
					vector<Point2f> cornersC;
					vector<Point2f> cornersG;
					if (!findChessboardCorners(color, patternSize, cornersC, 1+8)
						|| !findChessboardCorners(gray, patternSize, cornersG, 1+8))
						continue;
					cout << i << " " << j << endl;

					Mat drawImg = color;
					drawChessboardCorners(drawImg, patternSize, cornersC, 1);
					imshow("color_corner", drawImg);
					drawImg = gray;
					drawChessboardCorners(drawImg, patternSize, cornersG, 1);
					imshow("gray_corner", drawImg);

					ch = cvWaitKey(0);

					if (ch == 13) {
						cout << "save it..." << endl;
						finish = true;
						
						char prefix[100];
						sprintf(prefix, "data%d", saveIdx++);
						string prefixStr(prefix);
						
						imwrite(prefixStr+"_c.jpg", color);
						FileStorage fs(prefixStr+"_depthValue.yml", FileStorage::WRITE);
						fs << "mapped" << mapped;
						fs << "depth" << depth;
						fs.release();

						imwrite(prefixStr+"_g.jpg", gray);

						vector<Point3f> objpts;
						for (int a=0; a<j; ++a) {
							for (int b=0; b<i; ++b) {
								objpts.push_back(Point3f(b,a,0));
							}
						}
						objPts.push_back(objpts);
						imgCPts.push_back(cornersC);
						imgGPts.push_back(cornersG);

						break;
					}
				}
			}
		}
	}

	cout << "now calibration" << endl;
	Mat intriMatC, distortMatC, intriMatG, distortMatG;
	vector<Mat> rvecsC, tvecsC, rvecsG, tvecsG;
	calibrateCamera(objPts, imgCPts, Size(HEIGHT, WIDTH), intriMatC, distortMatC, rvecsC, tvecsC);
	ezPrintMat(intriMatC);
	ezPrintMat(distortMatC);
	calibrateCamera(objPts, imgGPts, Size(HEIGHT, WIDTH), intriMatG, distortMatG, rvecsG, tvecsG);
	ezPrintMat(intriMatG);
	ezPrintMat(distortMatG);

	Mat R,T,E,F;
	stereoCalibrate(objPts, imgCPts, imgGPts, intriMatC, distortMatC, intriMatG, distortMatG, 
		Size(HEIGHT,WIDTH), R, T, E, F);

	saveCaliResult("calibrate_color_result.yml", objPts, Size(HEIGHT,WIDTH), imgCPts, rvecsC, tvecsC,
		intriMatC, distortMatC);
	saveCaliResult("calibrate_gray_result.yml", objPts, Size(HEIGHT,WIDTH), imgGPts, rvecsG, tvecsG,
		intriMatG, distortMatG);
	
	FileStorage fs("stereo_result.yml", FileStorage::WRITE);
	fs << "R" << R;
	fs << "T" << T;
	fs << "E" << E;
	fs << "F" << F;
	
	fs.release();
}

class Color2GrayMapper {
public:
	Mat intriMatC, intriMatG;
	Mat distortMatC, distortMatG;
	vector<Mat> rvecsC, rvecsG, tvecsC, tvecsG;

	Mat R,T,R1,T1,R2,T2;
	double k1,k2,k3,p1,p2;

	void load(void) {
		FileStorage fs;
		FileNode vecs;

		fs.open("calibrate_color_result.yml", FileStorage::READ);
		fs["intriMat"] >> intriMatC;
		fs["distortMat"] >> distortMatC;
		vecs = fs["rvecs"];
		for (FileNodeIterator it = vecs.begin(); it != vecs.end(); ++it) {
			Mat tmp;
			(*it) >> tmp;
			rvecsC.push_back(tmp);
		}
		vecs = fs["tvecs"];
		for (FileNodeIterator it = vecs.begin(); it != vecs.end(); ++it) {
			Mat tmp;
			(*it) >> tmp;
			tvecsC.push_back(tmp);
		}
		fs.release();


		fs.open("calibrate_gray_result.yml", FileStorage::READ);
		fs["intriMat"] >> intriMatG;
		fs["distortMat"] >> distortMatG;
		vecs = fs["rvecs"];
		for (FileNodeIterator it = vecs.begin(); it != vecs.end(); ++it) {
			Mat tmp;
			(*it) >> tmp;
			rvecsG.push_back(tmp);
		}
		vecs = fs["tvecs"];
		for (FileNodeIterator it = vecs.begin(); it != vecs.end(); ++it) {
			Mat tmp;
			(*it) >> tmp;
			tvecsG.push_back(tmp);
		}
		fs.release();

		fs.open("stereo_result.yml", FileStorage::READ);
		fs["R"] >> R;
		fs["T"] >> T;
		fs.release();

	}

	void init(void) {
		load();

		for (int i=7; i<8; ++i) {
			Rodrigues(rvecsC[i], R1);
			T1 = tvecsC[i];
// 			ezPrintMat(R1);
// 			ezPrintMat(T1);

			Rodrigues(rvecsG[i], R2);
			T2 = tvecsG[i];
// 			ezPrintMat(R2);
// 			ezPrintMat(T2);

			Mat R = R2*R1.inv();
			Mat T = -R*T1 + T2;

			ezPrintMat(R);
			ezPrintMat(T);
		}
		ezPrintMat(R);
		ezPrintMat(T);


		k1 = distortMatG.at<double>(0);
		k2 = distortMatG.at<double>(1);
		k3 = distortMatG.at<double>(4);
		p1 = distortMatG.at<double>(2);
		p2 = distortMatG.at<double>(3);
	}

	Mat mapping(Mat P, short dval, double scale = 263) {
#ifdef VERBOSE
		ezPrintMat(P);
		ezPrintMat(intriMatC.inv()*P);
#endif
		P = intriMatC.inv()*P;
#ifdef VERBOSE
		ezPrintMat(dval/26.1 * P);
#endif
		P = dval/scale * P;

#ifdef VERBOSE
		ezPrintMat(R*P+T);
#endif
		P = R*P+T;

#ifdef VERBOSE
		ezPrintMat(P/P.at<double>(2));
#endif
		P = P/P.at<double>(2);

		double x,y,r;
		x = P.at<double>(0);
		y = P.at<double>(1);
		r = x*x+y*y;
		P.at<double>(0) = x*(1+k1*r+k2*r*r+k3*r*r*r) + 2*p1*x*y + p2*(r+2*x*x);
		P.at<double>(1) = y*(1+k1*r+k2*r*r+k3*r*r*r) + 2*p2*x*y + p1*(r+2*y*y);

#ifdef VERBOSE
		ezPrintMat(intriMatG*P);
#endif
		P = intriMatG*P;
		return P;
	}

};

void doMapping(void) {

	Color2GrayMapper mapper;
	mapper.init();

	char mode = '2';
	if (mode == '1') {
		Mat c0, d0, g0, dcMapped;
		Mat mapped = Mat::zeros(HEIGHT, WIDTH, CV_8UC3);
		FileStorage fs("data5_depthValue.yml", FileStorage::READ);
		fs["mapped"] >> dcMapped;
		fs["depth"] >> d0;
		fs.release();
		c0 = imread("data5_c1.jpg");
		g0 = imread("data5_g.jpg");

		imshow("c0", c0);
		imshow("g0", g0);
		imshow("d0", d0);
		imshow("mapped", mapped);
		cvWaitKey(0);

		const int MAXCOLOR = 20;
		int colorCnt = 0;
		Mat colorMat = createColor(MAXCOLOR);
		CvMat cvColorMat = colorMat;

		double base = 263;
		double offset = 2;
		double scales[] = {base, base-offset, base+offset};

		for (int test = 0; test < sizeof(scales)/sizeof(double); ++test) {
			for (int i=0; i<HEIGHT; ++i) {
				for (int j=0; j<WIDTH; ++j) {
					short dval = dcMapped.at<short>(i,j);
					if (dval <= 100 || dval >= 10000) continue;

					Mat P = (Mat_<double>(3,1) << j,i,1);

					P = mapper.mapping(P, dval, scales[test]);

					int px, py;
					px = cvRound(P.at<double>(0));
					py = cvRound(P.at<double>(1));

					if (0<=px && px<WIDTH && 0<=py && py<HEIGHT) {
						mapped.at<Vec3b>(py,px) = c0.at<Vec3b>(i,j);
					}
				}
			}
			
			imshow("c0", c0);
			imshow("g0", g0);
			imshow("mapped", mapped);
			mapped.setTo(0);

			cvWaitKey(0);
		}


	}
	else if (mode == '0') {
		Mat c0, d0, g0, dcMapped;
		Mat mapped = Mat::zeros(HEIGHT, WIDTH, CV_8UC3);
		FileStorage fs("test_data.yml", FileStorage::READ);
		fs["dcMapped"] >> dcMapped;
		fs["c0"] >> c0;
		fs["g0"] >> g0;
		fs["d0"] >> d0;
		fs.release();
		imshow("c0", c0);
		imshow("g0", g0);
		imshow("d0", d0);
		imshow("dcMapped", dcMapped);
		cvWaitKey(10);
		const int MAXCOLOR = 20;
		int colorCnt = 0;
		Mat colorMat = createColor(MAXCOLOR);
		CvMat cvColorMat = colorMat;

		for (int i=0; i<HEIGHT; ++i) {
			for (int j=0; j<WIDTH; ++j) {
				short dval = dcMapped.at<short>(i,j);
				if (dval <= 100 || dval >= 10000) continue;

				Mat P = (Mat_<double>(3,1) << j,i,1);

				P = mapper.mapping(P, dval);

				int px, py;
				px = cvRound(P.at<double>(0));
				py = cvRound(P.at<double>(1));
				
				if (0<=px && px<WIDTH && 0<=py && py<HEIGHT) {
					mapped.at<Vec3b>(py,px) = c0.at<Vec3b>(i,j);
// 					printf("point <%d,%d,%d> map to <%d,%d>\n", i,j,dval, py,px);
					//circle(g0, Point(px,py), 5, cvGet1D(&cvColorMat, colorCnt));
					//circle(c0, Point(j,i), 5, cvGet1D(&cvColorMat, colorCnt));
					//colorCnt = (colorCnt+1)%MAXCOLOR;
// 					imshow("debugG0", g0);
// 					imshow("debugC0", c0);
// 					imshow("debugMapped", mapped);
// 					cvWaitKey(0);
				}
			}
		}

		imshow("g0", g0);
		imshow("mapped", mapped);
		while (1)
		cvWaitKey(0);

	}
	else {
		//reading from device
		Mat mapped(HEIGHT, WIDTH, CV_8UC3);
		Mat dcMapped(HEIGHT, WIDTH, CV_16S);
		KinectDepthColorWrapper kdcWrapper;
		kdcWrapper.init(WIDTH, HEIGHT);
		VideoCapture capture(CV_CAP_DSHOW+1);
		capture.set(CV_CAP_PROP_FRAME_WIDTH, WIDTH);
		capture.set(CV_CAP_PROP_FRAME_HEIGHT, HEIGHT);

		long *colorCoordinates = new long[WIDTH*HEIGHT*2];
		Mat c0, d0, g0;
		CvMat *depthRes = NULL, *colorRes = NULL;

		while (true) {
			for (int i=0; i<HEIGHT; ++i) {
				for (int j=0; j<WIDTH; ++j) {
					colorCoordinates[2*(i*WIDTH+j)]=j;
					colorCoordinates[2*(i*WIDTH+j)+1]=i;
					memset(mapped.ptr<uchar>(i,j), 0, sizeof(uchar)*3);
					dcMapped.at<short>(i,j) = 0;
				}
			}
			kdcWrapper.retrieve(depthRes, colorRes, colorCoordinates);
			capture.read(g0);
			flip(g0, g0, 1);

			d0 = Mat(depthRes);
			c0 = Mat(colorRes);

			imshow("depth", d0);
			imshow("color", c0);
			imshow("gray", g0);
			char ch = cvWaitKey(0);
			if (ch != 13) continue;

			for (int i=0; i<HEIGHT; ++i) {
				for (int j=0; j<WIDTH; ++j) {
					int d2cX, d2cY;
					d2cX = colorCoordinates[2*(i*WIDTH+j)];
					d2cY = colorCoordinates[2*(i*WIDTH+j)+1];
					if (0<=d2cX && d2cX<WIDTH && 0<=d2cY && d2cY<HEIGHT && d0.at<short>(i,j) > 100) {
						short dval = d0.at<short>(i,j);
						dcMapped.at<short>(d2cY,d2cX) = dval;

						Mat P = (Mat_<double>(3,1) << d2cX,d2cY,1);

						P = mapper.mapping(P, dval);

						int px, py;
						px = cvRound(P.at<double>(0));
						py = cvRound(P.at<double>(1));
						if (0<=px && px<WIDTH && 0<=py && py<HEIGHT) {
							memcpy(mapped.ptr<uchar>(py,px), c0.ptr<uchar>(d2cY, d2cX), sizeof(uchar)*3);
						}
					}
				}
			}

			FileStorage fs("test_data.yml", FileStorage::WRITE);
			fs << "dcMapped" << dcMapped;
			fs << "c0" << c0;
			fs << "g0" << g0;
			fs << "d0" << d0;
			fs.release();

			imshow("gray", g0);
			imshow("mapped", dcMapped);
			ch = cvWaitKey(0);
			if (ch == 27) break;
		}
	}
}


int main(void)
{
	char mode = '0';
	cout << "choose mode: (c)color or (g)gray or (0) load both" << endl;
//	cin >> mode;
	cout << "chose mode " << mode << endl;

	if (mode == 'b') {
		caliBoth();
	}
	else if (mode == 'd' || mode == 'c') {
		caliColor();
	}
	else if (mode == 'g') {
		caliGray();
	}
	else if (mode == '0') {
		doMapping();
	}

	return 0;
}